{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example notebook - pandas [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/spark/world_of_warcraft/world_of_warcraft__pandas.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/spark/world_of_warcraft/world_of_warcraft__pandas.ipynb)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-22T22:27:16.952106Z",
     "start_time": "2023-05-22T22:27:15.116241Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:'PYARROW_IGNORE_TIMEZONE' environment variable was not set. It is required to set this environment variable to '1' in both driver and executor sides if you use pyarrow>=2.0.0. pandas-on-Spark will set it for you but it does not work if there is a Spark context already launched.\n"
     ]
    }
   ],
   "source": [
    "# Cell 1 - import the things you need\n",
    "import logging\n",
    "import sys\n",
    "import importlib\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from hamilton import ad_hoc_utils, driver\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-22T22:27:16.982107Z",
     "start_time": "2023-05-22T22:27:16.954791Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Cell 2 - import modules to create part of the DAG from\n",
    "# We use the autoreload extension that comes with ipython to automatically reload modules when\n",
    "# the code in them changes.\n",
    "\n",
    "# import the jupyter extension\n",
    "%load_ext autoreload\n",
    "# set it to only reload the modules imported\n",
    "%autoreload 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hamilton.telemetry:Note: Hamilton collects completely anonymous data about usage. This will help us improve Hamilton over time. See https://github.com/apache/hamilton#usage-analytics--data-privacy for details.\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"496pt\" height=\"404pt\"\n",
       " viewBox=\"0.00 0.00 496.24 404.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 400)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-400 492.24,-400 492.24,4 -4,4\"/>\n",
       "<!-- zone -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>zone</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"131.69\" cy=\"-378\" rx=\"50.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"131.69\" y=\"-374.3\" font-family=\"Times,serif\" font-size=\"14.00\">Input: zone</text>\n",
       "</g>\n",
       "<!-- darkshore_flag -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>darkshore_flag</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"63.69\" cy=\"-306\" rx=\"63.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"63.69\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">darkshore_flag</text>\n",
       "</g>\n",
       "<!-- zone&#45;&gt;darkshore_flag -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>zone&#45;&gt;darkshore_flag</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M115.92,-360.76C107.55,-352.14 97.08,-341.36 87.71,-331.73\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"90.39,-329.45 80.91,-324.72 85.36,-334.33 90.39,-329.45\"/>\n",
       "</g>\n",
       "<!-- durotar_flag -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>durotar_flag</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"199.69\" cy=\"-306\" rx=\"53.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"199.69\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">durotar_flag</text>\n",
       "</g>\n",
       "<!-- zone&#45;&gt;durotar_flag -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>zone&#45;&gt;durotar_flag</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M147.47,-360.76C155.84,-352.14 166.31,-341.36 175.67,-331.73\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"178.02,-334.33 182.48,-324.72 173,-329.45 178.02,-334.33\"/>\n",
       "</g>\n",
       "<!-- durotar_count -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>durotar_count</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"152.69\" cy=\"-90\" rx=\"59.59\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"152.69\" y=\"-86.3\" font-family=\"Times,serif\" font-size=\"14.00\">durotar_count</text>\n",
       "</g>\n",
       "<!-- durotar_likelihood -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>durotar_likelihood</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"195.69\" cy=\"-18\" rx=\"76.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"195.69\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">durotar_likelihood</text>\n",
       "</g>\n",
       "<!-- durotar_count&#45;&gt;durotar_likelihood -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>durotar_count&#45;&gt;durotar_likelihood</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M163.1,-72.05C167.89,-64.26 173.69,-54.82 179.06,-46.08\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"182.04,-47.92 184.29,-37.56 176.07,-44.25 182.04,-47.92\"/>\n",
       "</g>\n",
       "<!-- darkshore_likelihood -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>darkshore_likelihood</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"396.69\" cy=\"-18\" rx=\"85.29\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"396.69\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">darkshore_likelihood</text>\n",
       "</g>\n",
       "<!-- darkshore_count -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>darkshore_count</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"418.69\" cy=\"-90\" rx=\"69.59\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"418.69\" y=\"-86.3\" font-family=\"Times,serif\" font-size=\"14.00\">darkshore_count</text>\n",
       "</g>\n",
       "<!-- darkshore_count&#45;&gt;darkshore_likelihood -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>darkshore_count&#45;&gt;darkshore_likelihood</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M413.26,-71.7C410.91,-64.24 408.11,-55.32 405.48,-46.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"408.86,-46.04 402.52,-37.55 402.18,-48.14 408.86,-46.04\"/>\n",
       "</g>\n",
       "<!-- zone_counts -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>zone_counts</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"280.69\" cy=\"-162\" rx=\"54.69\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"280.69\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">zone_counts</text>\n",
       "</g>\n",
       "<!-- zone_counts&#45;&gt;durotar_count -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>zone_counts&#45;&gt;durotar_count</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M253.24,-145.98C234.97,-135.99 210.73,-122.74 190.63,-111.75\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"192.33,-108.69 181.88,-106.96 188.97,-114.83 192.33,-108.69\"/>\n",
       "</g>\n",
       "<!-- zone_counts&#45;&gt;darkshore_count -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>zone_counts&#45;&gt;darkshore_count</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M309.62,-146.33C329.55,-136.22 356.29,-122.65 378.31,-111.49\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"379.68,-114.71 387.02,-107.07 376.52,-108.47 379.68,-114.71\"/>\n",
       "</g>\n",
       "<!-- total_count -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>total_count</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"280.69\" cy=\"-90\" rx=\"50.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"280.69\" y=\"-86.3\" font-family=\"Times,serif\" font-size=\"14.00\">total_count</text>\n",
       "</g>\n",
       "<!-- zone_counts&#45;&gt;total_count -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>zone_counts&#45;&gt;total_count</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M280.69,-143.7C280.69,-136.41 280.69,-127.73 280.69,-119.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"284.19,-119.62 280.69,-109.62 277.19,-119.62 284.19,-119.62\"/>\n",
       "</g>\n",
       "<!-- avatarId -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>avatarId</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"334.69\" cy=\"-306\" rx=\"63.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"334.69\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">Input: avatarId</text>\n",
       "</g>\n",
       "<!-- with_flags -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>with_flags</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"199.69\" cy=\"-234\" rx=\"48.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"199.69\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">with_flags</text>\n",
       "</g>\n",
       "<!-- avatarId&#45;&gt;with_flags -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>avatarId&#45;&gt;with_flags</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M305.07,-289.64C285.06,-279.26 258.48,-265.48 237.03,-254.36\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"238.79,-251.33 228.3,-249.83 235.57,-257.54 238.79,-251.33\"/>\n",
       "</g>\n",
       "<!-- aggregation_level -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>aggregation_level</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"362.69\" cy=\"-234\" rx=\"96.68\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"362.69\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">Input: aggregation_level</text>\n",
       "</g>\n",
       "<!-- aggregation_level&#45;&gt;zone_counts -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>aggregation_level&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M342.84,-216.05C332.52,-207.24 319.74,-196.33 308.45,-186.69\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"310.83,-184.12 300.95,-180.29 306.28,-189.44 310.83,-184.12\"/>\n",
       "</g>\n",
       "<!-- with_flags&#45;&gt;zone_counts -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>with_flags&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M218.07,-217.12C228.49,-208.12 241.72,-196.68 253.33,-186.65\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"255.4,-189.48 260.68,-180.3 250.82,-184.19 255.4,-189.48\"/>\n",
       "</g>\n",
       "<!-- total_count&#45;&gt;darkshore_likelihood -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>total_count&#45;&gt;darkshore_likelihood</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M305.58,-73.98C321.42,-64.42 342.21,-51.88 359.94,-41.18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"361.37,-44.4 368.12,-36.24 357.75,-38.41 361.37,-44.4\"/>\n",
       "</g>\n",
       "<!-- total_count&#45;&gt;durotar_likelihood -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>total_count&#45;&gt;durotar_likelihood</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M261.41,-73.12C250.63,-64.24 236.99,-53.01 224.93,-43.08\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"227.2,-40.41 217.26,-36.76 222.75,-45.82 227.2,-40.41\"/>\n",
       "</g>\n",
       "<!-- darkshore_flag&#45;&gt;with_flags -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>darkshore_flag&#45;&gt;with_flags</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M93.54,-289.64C113.7,-279.26 140.47,-265.48 162.09,-254.36\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"163.6,-257.52 170.89,-249.83 160.39,-251.29 163.6,-257.52\"/>\n",
       "</g>\n",
       "<!-- durotar_flag&#45;&gt;with_flags -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>durotar_flag&#45;&gt;with_flags</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M199.69,-287.7C199.69,-280.41 199.69,-271.73 199.69,-263.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"203.19,-263.62 199.69,-253.62 196.19,-263.62 203.19,-263.62\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x1236f9f70>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_df = pd.read_parquet(\"data/wow.parquet\")\n",
    "\n",
    "aggregation_level = \"avatarId\"\n",
    "zone_features__pandas = importlib.import_module(\"zone_features__pandas\")\n",
    "zone_features__common = importlib.import_module(\"zone_features__common\")\n",
    "\n",
    "dr = driver.Driver({}, zone_features__common, zone_features__pandas)  # can pass in multiple modules\n",
    "dr.display_all_functions(None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          total_count  darkshore_count  darkshore_likelihood\n",
      "avatarId                                                    \n",
      "0                  76                0                   0.0\n",
      "1                  19                0                   0.0\n",
      "2                 192                0                   0.0\n",
      "3                 109                0                   0.0\n",
      "4                  76                0                   0.0\n",
      "...               ...              ...                   ...\n",
      "91060               1                0                   0.0\n",
      "91061               2                0                   0.0\n",
      "91062               5                0                   0.0\n",
      "91063               2                0                   0.0\n",
      "91064               1                0                   0.0\n",
      "\n",
      "[91064 rows x 3 columns]\n"
     ]
    }
   ],
   "source": [
    "# we need to specify what we want in the final dataframe.\n",
    "output_columns = [\n",
    "    \"total_count\",\n",
    "    \"darkshore_count\",\n",
    "    \"darkshore_likelihood\",\n",
    "]\n",
    "\n",
    "# let's create the dataframe!\n",
    "# if you only did `pip install sf-hamilton` earlier:\n",
    "df = dr.execute(output_columns, inputs={**input_df.to_dict(orient='series'), \"aggregation_level\": aggregation_level})\n",
    "# else if you did `pip install \"sf-hamilton[visualization]\"` earlier:\n",
    "# dr.visualize_execution(output_columns, './my-dag.dot', {})\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
